package com.wt.admin.util;

import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.bytedeco.ffmpeg.global.avcodec;
import org.bytedeco.ffmpeg.global.avutil;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.*;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.*;
import java.util.function.Function;


@Slf4j
public class VideoUtil {


    /**
     * 处理视频，支持调整分辨率
     */
    public static void video(String inputPath, String outputPath, int targetWidth, int targetHeight, Function<BufferedImage, Graphics2D> fun)
            throws IOException, FrameGrabber.Exception, FrameRecorder.Exception {
        // 初始化视频抓取器
        try (FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(inputPath)) {
            grabber.start();
            // 调整分辨率
            grabber.setImageWidth(targetWidth);
            grabber.setImageHeight(targetHeight);
            // 获取视频的帧率、宽度和高度
            int width = grabber.getImageWidth();
            int height = grabber.getImageHeight();
            // 初始化视频录制器
            try (FFmpegFrameRecorder recorder = new FFmpegFrameRecorder(outputPath, width, height)) {
                recorder.setVideoCodec(avcodec.AV_CODEC_ID_H264);
                recorder.setFormat("mp4");
                recorder.setFrameRate(grabber.getFrameRate());
                recorder.setVideoBitrate(5_000_000); // 设置高比特率以保持画质
                // 设置音频参数
                recorder.setAudioChannels(grabber.getAudioChannels());
                recorder.setSampleRate(grabber.getSampleRate());
                recorder.setAudioCodec(avcodec.AV_CODEC_ID_AAC);
                recorder.start();
                Frame frame;
                int index = 0;
                // 使用 try-with-resources 管理 Java2DFrameConverter
                try (Java2DFrameConverter converter = new Java2DFrameConverter()) {
                    while ((frame = grabber.grab()) != null) {
                        // 过滤非视频帧 并 只处理第 10 帧  && frameIndex == 10
                        if (frame.image != null && index % 5 == 0) {
                            index ++;
                            BufferedImage image = converter.getBufferedImage(frame);
                            Graphics2D g2d = fun.apply(image);
                            g2d.dispose();
                            // 将修改后的图片转换回 Frame
                            frame = converter.convert(image);
                        }
                        // 将帧写入输出视频
                        recorder.record(frame);
                    }
                }
            }
        }
    }

    /**
     * 处理视频并返回实时视频流
     *
     * @param inputPath    输入视频路径
     * @param targetWidth  目标宽度
     * @param targetHeight 目标高度
     * @param outputPath          每一帧的消费者（用于实时处理和返回视频流）
     * @throws IOException            如果发生 I/O 错误
     * @throws FrameGrabber.Exception 如果抓取帧时发生错误
     */
    public static void videoToStream(String inputPath, String outputPath, int targetWidth, int targetHeight, Function<BufferedImage, Graphics2D> fun)
            throws IOException, FrameGrabber.Exception {
        // 初始化视频抓取器
        log.debug("视频识别开始");
        try (FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(inputPath)) {
            grabber.start();
            // 获取视频的帧率（每秒帧数）
            int frameRate = (int) grabber.getFrameRate();
            log.debug("视频每秒帧数 {}", frameRate);
            // 调整分辨率
            grabber.setImageWidth(targetWidth);
            grabber.setImageHeight(targetHeight);

            // 获取视频的帧率、宽度和高度
            int width = grabber.getImageWidth();
            int height = grabber.getImageHeight();
            // 初始化视频录制器
            try (FFmpegFrameRecorder recorder = new FFmpegFrameRecorder(outputPath, width, height)) {
                // 设置 HLS 参数，减少延迟
                recorder.setOption("hls_time", "3"); // 每个切片时长 2 秒
                recorder.setOption("hls_list_size", "3"); // 播放列表中保留 3 个切片
                recorder.setOption("hls_flags", "delete_segments+append_list");
//                recorder.setOption("hls_segment_filename", "D:\\download\\server\\file\\test\\video_%04d.ts");
                recorder.setOption("hls_allow_cache", "0");    // 禁用客户端缓存
                recorder.setOption("hls_playlist_type", "event"); // 适用于实时流类型
                recorder.setVideoCodec(avcodec.AV_CODEC_ID_H264);
                // 添加关键帧间隔配置（GOP大小）
                recorder.setVideoOption("g", String.valueOf((int)(grabber.getFrameRate() * 2))); // 2秒一个GOP
                recorder.setVideoOption("profile:v", "baseline"); // 兼容性更好的编码配置
                recorder.setFormat("hls");
                recorder.setFrameRate(grabber.getFrameRate());
                recorder.setVideoBitrate(5_000_000); // 设置高比特率以保持画质
                recorder.setPixelFormat(avutil.AV_PIX_FMT_YUV420P);

                // 设置音频参数
                if (grabber.getAudioChannels() > 0) {
                    recorder.setAudioCodec(avcodec.AV_CODEC_ID_AAC);
                    recorder.setAudioChannels(grabber.getAudioChannels());
                    recorder.setSampleRate(grabber.getSampleRate());
                }
                recorder.start();
                Frame frame;
                int index = 0;
                // 计算10%的帧数并向上取整
                int framesToProcess = (int) Math.ceil(frameRate * 0.1);
                int interval = (int) Math.ceil((double) frameRate / framesToProcess);
                // 使用 try-with-resources 管理 Java2DFrameConverter
                try (Java2DFrameConverter converter = new Java2DFrameConverter()) {
                    while ((frame = grabber.grab()) != null) {
                        // 过滤非视频帧
                        if(frame.image == null){
                            recorder.record(frame);
                            continue;
                        }
                        // 处理图片帧
                        if (index % interval == 0) {
                            BufferedImage image = converter.getBufferedImage(frame);
                            Graphics2D g2d = fun.apply(image);
                            g2d.dispose();
                            // 将修改后的图片转换回 Frame
                            frame = converter.convert(image);
                        }
                        index ++;
                        // 将帧写入输出视频
                        recorder.record(frame);
                    }
                }
            }
        }
        log.debug("视频识别结束");
    }


    /**
     * 处理视频并输出新视频，支持调整分辨率或保持原始分辨率
     *
     * @param inputPath    输入视频路径
     * @param outputPath   输出视频路径
     * @param targetWidth  目标宽度（为0时保持原始宽度）
     * @param targetHeight 目标高度（为0时保持原始高度）
     * @throws IOException             如果发生 I/O 错误
     * @throws FrameGrabber.Exception  如果抓取帧时发生错误
     * @throws FrameRecorder.Exception 如果录制帧时发生错误
     */
    public static void resizeVideoWH(String inputPath, String outputPath, int targetWidth, int targetHeight)
            throws IOException, FrameGrabber.Exception, FrameRecorder.Exception {
        // 初始化视频抓取器
        try (FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(inputPath)) {
            grabber.start();
            // 调整分辨率
            grabber.setImageWidth(targetWidth);
            grabber.setImageHeight(targetHeight);
            // 获取视频的帧率
            double frameRate = grabber.getFrameRate() > 0 ? grabber.getFrameRate() : 30; // 默认帧率为 30

            // 初始化视频录制器
            try (FFmpegFrameRecorder recorder = new FFmpegFrameRecorder(outputPath, targetWidth, targetHeight)) {
                recorder.setVideoCodec(avcodec.AV_CODEC_ID_H264);
                recorder.setFormat(getOutputFormat(outputPath));
                recorder.setFrameRate(frameRate);
                recorder.setVideoBitrate(5_000_000); // 设置高比特率以保持画质

                // 设置音频参数
                recorder.setAudioChannels(grabber.getAudioChannels());
                recorder.setSampleRate(grabber.getSampleRate());
                recorder.setAudioCodec(avcodec.AV_CODEC_ID_AAC);

                recorder.start();

                Frame frame;
                while ((frame = grabber.grab()) != null) {
                    // 将帧写入输出视频
                    recorder.record(frame);
                }
            } // 自动释放 recorder 资源
        } // 自动释放 grabber 资源
    }

    public static void resizeVideoWH(InputStream stream, String outputPath, int targetWidth, int targetHeight)
            throws IOException, FrameGrabber.Exception, FrameRecorder.Exception {
        // 初始化视频抓取器
        try (FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(stream)) {
            grabber.start();
            // 调整分辨率
            grabber.setImageWidth(targetWidth);
            grabber.setImageHeight(targetHeight);
            // 获取视频的帧率
            double frameRate = grabber.getFrameRate() > 0 ? grabber.getFrameRate() : 30; // 默认帧率为 30

            // 初始化视频录制器
            try (FFmpegFrameRecorder recorder = new FFmpegFrameRecorder(outputPath, targetWidth, targetHeight)) {
                recorder.setVideoCodec(avcodec.AV_CODEC_ID_H264);
                recorder.setFormat(getOutputFormat(outputPath));
                recorder.setFrameRate(frameRate);
                recorder.setVideoBitrate(5_000_000); // 设置高比特率以保持画质

                // 设置音频参数
                recorder.setAudioChannels(grabber.getAudioChannels());
                recorder.setSampleRate(grabber.getSampleRate());
                recorder.setAudioCodec(avcodec.AV_CODEC_ID_AAC);

                recorder.start();

                Frame frame;
                while ((frame = grabber.grab()) != null) {
                    // 将帧写入输出视频
                    recorder.record(frame);
                }
            } // 自动释放 recorder 资源
        } // 自动释放 grabber 资源
    }


    private static String getOutputFormat(String outputPath) {
        // 校验 outputPath 是否包含扩展名
        if (outputPath == null || outputPath.lastIndexOf(".") == -1) {
            return "mp4"; // 默认格式为 mp4
        }
        String extension = outputPath.substring(outputPath.lastIndexOf(".") + 1);
        return switch (extension.toLowerCase()) {
            case "mp4" -> "mp4";
            case "avi" -> "avi";
            case "mkv" -> "matroska";
            default -> "mp4"; // 默认格式为 mp4
        };
    }

    @SneakyThrows
    public static InputStream bufferedImageToInputStream(BufferedImage image) {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        ImageIO.write(image, "jpg", outputStream);
        return new ByteArrayInputStream(outputStream.toByteArray());
    }

}
